// my_gemv_kernel.cu
#include <cuda.h>
#include <cuda_runtime.h>
#include <torch/extension.h>

template <typename scalar_t>
__global__ void my_gemv_kernel(const scalar_t *__restrict__ A, // (M, K)
                               const scalar_t *__restrict__ x, // (K,)
                               scalar_t *__restrict__ y,       // (M,)
                               int M, int K) {
  int row = blockIdx.x * blockDim.x + threadIdx.x;
  if (row < M) {
    scalar_t sum = 0;
    for (int k = 0; k < K; ++k)
      sum += A[row * K + k] * x[k];
    y[row] = sum;
  }
}

torch::Tensor my_gemv_cuda(const torch::Tensor &A, const torch::Tensor &x) {
  TORCH_CHECK(A.device().is_cuda(), "A must be CUDA");
  TORCH_CHECK(x.device().is_cuda(), "x must be CUDA");
  TORCH_CHECK(A.dim() == 2 && x.dim() == 1);
  TORCH_CHECK(A.size(1) == x.size(0));

  int M = A.size(0);
  int K = A.size(1);
  auto y = torch::zeros({M}, A.options());

  const int threads = 256;
  const int blocks = (M + threads - 1) / threads;

  AT_DISPATCH_FLOATING_TYPES(A.scalar_type(), "my_gemv_cuda", ([&] {
                               my_gemv_kernel<scalar_t><<<blocks, threads>>>(
                                   A.data_ptr<scalar_t>(),
                                   x.data_ptr<scalar_t>(),
                                   y.data_ptr<scalar_t>(), M, K);
                             }));
  return y;
}